from urllib.request import *
from urllib.parse import urlencode

from os.path import join

import config
config = config.Config()

def extract_problems(xml):
    for node in xml.findall("*/problem"):
        yield (node.get("index"),
               node.get("name"),
               node.get("url") + "?type=linux"
               ) 

def download_problem(dst_dir, id, url):
    print("Download problem %s by url %s" % (id, url))
    zip_content = join(dst_dir, "prob_%s.zip") % id
    urlretrieve (
            url,
            zip_content,
            data = urlencode(config.access))

    from zipfile import ZipFile
    ZipFile(zip_content).extractall(join(dst_dir, id))

def find_statement(doc):
    long_name = None
    latex_path = None
    for s in doc.findall("statements/statement"):
        if s.get("language") == config.language:
            #support of old packages
            title = s.get("title")
            if title is not None:
                long_name = '"' + title + '"'
            latex_path = s.get("path")
            statement_found = True

    #new packages
    for s in doc.findall("names/name"):
        if s.get("language") == config.language:
            title = s.get("value")
            if title is not None:
                long_name = '"' + title + '"'
        
    return long_name, latex_path


def extract_problem_meta(dir):
    import xml.etree.ElementTree as ET
    doc = ET.ElementTree(file = join(dir, "problem.xml"))
    res = {}
    root = doc.getroot()
    short_name = root.get("name")
    if short_name is None:
        short_name = root.get("short-name")

    if short_name is not None:
        res["short_name"] = '"' + short_name + '"'
    path = {}
    test_set = doc.findall("judging/testset")
    for ts in test_set:
        if ts.get("name") == "tests":
            test_set = ts
            break

    assert(not isinstance(test_set, list))

    res["time_limit_millis"] = test_set.findtext("time-limit")
    res["max_vm_size"] = test_set.findtext("memory-limit")

    judging = doc.find("judging")
    inf = judging.get("input-file")
    res["input_file"] = '"' + inf + '"'
    res["use_stdin"] = "1" if len(inf) == 0 else "0"

    out = judging.get("output-file")
    res["output_file"] = '"' + out + '"'
    res["use_stdout"] = "1" if len(out) == 0 else "0"

    path["checker"] = join(dir, doc.find("assets/checker/source").get("path"))

    path["resources"] = []
    for r in doc.findall("files/resources/file"):
        path["resources"].append(join(dir, r.get("path")))

    long_name, latex_path = find_statement(doc)

    if not long_name or not latex_path:
        raise Exception(
            "couldn't found statement for problem %s on language %s"
            % (dir, config.language)
        )

    res["long_name"] = long_name
    path["latex_path"] = join(dir, latex_path)

    return res, path

def install_problem(tmp_dir, ejudge_dir, prob, paths):
    from os.path import exists, basename
    from os import makedirs
    from shutil import move, copy, copytree
    
    if config.separate_dir_for_problem:
        ejudge_prob_dir = join(ejudge_dir, "problems", prob["short_name"][1:-1])
    else:
        checker_id = prob["id"]
        prob["short_name"] = '"%s"' % checker_id
        
        ejudge_test_dir = join(ejudge_dir, "tests", checker_id)
        ejudge_checker_dir = join(ejudge_dir, "checkers")

        ejudge_checker_binary = join(ejudge_checker_dir, "check_%s" % checker_id) 
        ejudge_checker_srcs = [join(ejudge_checker_dir, "check_%s.%s" % (checker_id, ext)) for ext in ["cpp", "pas", "dpr"]]  

    src_prob_dir = join(tmp_dir, prob["id"])

    ejudge_backup = join(tmp_dir, "ejudge")
    if not exists(ejudge_backup):
        makedirs(ejudge_backup)

    if config.separate_dir_for_problem:
        if exists(ejudge_prob_dir):
            print("moving problem %s to backup %s" % (prob["id"], ejudge_backup))
            move(ejudge_prob_dir, ejudge_backup)
    else:
        if exists(ejudge_test_dir):
            print("moving problem tests %s to backup %s" % (ejudge_test_dir, ejudge_backup))
            move(ejudge_test_dir, ejudge_backup)
        for check_file in ejudge_checker_srcs + [ejudge_checker_binary]:
            if exists(check_file):
                print("moving checker %s to backup %s" % (check_file, ejudge_backup))
                move(check_file, ejudge_backup)

    if config.separate_dir_for_problem:
         makedirs(ejudge_prob_dir)
         for directory in ["tests", "statements"]:
            copytree(join(src_prob_dir, directory), join(ejudge_prob_dir, directory))
    else:
        if not exists(ejudge_checker_dir): makedirs(ejudge_checker_dir)
        copytree(join(src_prob_dir, "tests"), ejudge_test_dir)

    checker_ext = basename(paths["checker"]).split(".")[-1]
    if config.separate_dir_for_problem:
        checker_file = join(ejudge_prob_dir, "check." + checker_ext)
    else:
        checker_file = join(ejudge_checker_dir, "check_%s.%s" % (checker_id, checker_ext))
    print("Moving checker %s to %s " % (paths["checker"], checker_file))
    copy(paths["checker"], checker_file)

    if checker_ext == "cpp":
        print("Adding define EJUDGE to checker")
        lines = [
        """/* Autogenerated code, forces testlib to return exit codes for EJUDGE. */\n""",
        """#define EJUDGE\n""",
        "\n"]
        with open(checker_file) as f:
            lines += list(f.readlines())
        with open(checker_file, "w") as f:
            f.writelines(lines)

    for cfile in paths["resources"]:
        if basename(cfile) not in config.ignore_resources:
            if checker_ext != "cpp" and basename(cfile) in config.cpp_only_resources:
                continue
            print("Copy recource file %s" % cfile)
            copy(cfile, ejudge_prob_dir if config.separate_dir_for_problem else ejudge_checker_dir)
